package linq.mlinq;

import haxe.macro.Compiler;
import haxe.macro.Context;
import haxe.macro.Expr;
import tink.macro.tools.AST;
using tink.macro.tools.MacroTools;

class Enumerable<T> {
	public var elems:Array<T>;

	public function new () {
		elems = new Array<T> ();
	}

	@:macro public function Where (enumerable:Expr, predicate:Expr):Expr {
		switch (enumerable.expr) {
			case EBlock(block):
				var chain = chains.get (Context.signature (block[block.length-1]));
				var inner = AST.build (out.elems.push (elem));
				var ast = AST.build (if ($predicate (elem)) $inner);
				chain.inner.expr = ast.expr;
				chain.inner = inner;
				return chain.base;
			default: trace ("Nothing");
		}
		return enumerable;
	}

	@:macro public function Count (enumerable:Expr, ?predicate:Expr):Expr {
		switch (enumerable.expr) {
			case EBlock(block):
				var chain = chains.get (Context.signature (block[block.length-1]));
				chain.out.expr = tCount ().expr;
				var inner = AST.build (++out);
				var ast = if (isNull (predicate))
					AST.build ($inner);
				else
					AST.build (if ($predicate (elem)) $inner);				
				chain.inner.expr = ast.expr;
				chain.inner = inner;
				return chain.base;
			default: trace ("Nothing");
		}
		return enumerable;
	}

	@:macro public function Take (enumerable:Expr, num:Expr):Expr {
		switch (enumerable.expr) {
			case EBlock(block):
				var chain = chains.get (Context.signature (block[block.length-1]));
				var inner = AST.build (out.elems.push (elem));
				var ast = AST.build (if (index < $num) $inner else break);		
				chain.inner.expr = ast.expr;
				chain.inner = inner;
				return chain.base;
			default: trace ("Nothing");
		}
		return enumerable;
	}

	@:macro public function TakeWhile (enumerable:Expr, predicate:Expr):Expr {
		switch (enumerable.expr) {
			case EBlock(block):
				var chain = chains.get (Context.signature (block[block.length-1]));
				var inner = AST.build (out.elems.push (elem));
				var ast = AST.build (if ($predicate (elem)) $inner else break);		
				chain.inner.expr = ast.expr;
				chain.inner = inner;
				return chain.base;
			default: trace ("Nothing");
		}
		return enumerable;
	}

	@:macro public function Select (enumerable:Expr, predicate:Expr):Expr {
		switch (enumerable.expr) {
			case EBlock(block):
				var chain = chains.get (Context.signature (block[block.length-1]));
				chain.out.expr = tEnumerable (predicate);
				var inner = AST.build (out.elems.push (elem));
				var ast = AST.build ({var elem = $predicate (elem); $inner;});
				chain.inner.expr = ast.expr;
				chain.inner = inner;
				return chain.base;
			default: trace ("Nothing");
		}
		return enumerable;
	}

	@:macro public function Chain (enumerable:Expr):Expr {
		var inner = {expr:EBlock ([]), pos:Context.currentPos ()};



		var out = {expr:null, pos:enumerable.pos};
		out.expr = tEnumerable (enumerable);
		var ast = AST.build ({
			$out;
			var index:Int = 0;
			for (elem in $enumerable.elems) {
				$inner;
				++index;
			}
			$enumerable;
		});

		chains.set (Context.signature (enumerable),{base:ast, inner:inner, out:out});

		return ast;
	}

	#if macro
	static function tEnumerable (elem:Expr) {
		var type = Context.getType ("Enumerable");
		switch (Context.typeof (elem)) {
			case TInst (a,b):
				switch (type) {
					case TInst (c,d):
						d.pop ();
						d.push (b[0]);
					default:
				}
			default:
		}
		var type = type.toComplex ();
		var arr = "out".define (AST.build (new Enumerable ()), type, elem.pos);
		return arr.expr;
	}

	static function tArray (elem:Expr) {
		var type = Context.getType ("Enumerable");
		switch (type) {
			case TInst (c,d):
				d.pop ();
				d.push (Context.typeof (elem));
			default:
		}
		trace ("tArray Type: " + type);
		var type = type.toComplex ();
		var arr = "out".define (AST.build (new Enumerable ()), type, elem.pos);
		return arr.expr;
	}

	static function tCount () {
		return "out".define (AST.build (0), Context.currentPos ());
	}

	static function isNull (expr:Expr):Bool {
		switch (expr.expr) {
			case EConst (c):
				switch (c) {
					case CIdent (s):
						return s == "null";
					default:
				}
			default:
		}
		return false;
	}
	#end

	static var chains:Hash<{base:Expr, inner:Expr, out:Expr}> = new Hash ();
}